import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


# 1. 数据加载（保持简洁，移除数据增强）
def load_data_mnist(batch_size):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_data = datasets.MNIST('./data', train=False, transform=transform)

    return (
        DataLoader(train_data, batch_size, shuffle=True, num_workers=2),
        DataLoader(test_data, batch_size, shuffle=False, num_workers=2)
    )


# 2. 带正则化的简洁模型
class RegularizedNN(nn.Module):
    def __init__(self, dropout_rate=0.5):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Dropout(dropout_rate),  # 核心正则化
            nn.Linear(256, 10)
        )

        # 初始化权重
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.01)
                nn.init.zeros_(m.bias)

    def forward(self, X):
        return self.net(X)


# 3. 训练函数（保持核心正则化）
def train_model(net, train_iter, test_iter, num_epochs=10, lr=0.1, weight_decay=1e-4):
    # 使用SGD优化器并添加L2正则化（weight_decay）
    optimizer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=weight_decay)
    criterion = nn.CrossEntropyLoss()

    train_loss_list, train_acc_list, test_acc_list = [], [], []

    for epoch in range(num_epochs):
        net.train()
        total_loss, total_correct = 0, 0

        for X, y in train_iter:
            optimizer.zero_grad()
            output = net(X)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            total_correct += (output.argmax(1) == y).sum().item()

        # 计算指标
        train_loss = total_loss / len(train_iter)
        train_acc = total_correct / len(train_iter.dataset)

        # 测试集评估
        net.eval()
        test_correct = 0
        with torch.no_grad():
            for X, y in test_iter:
                test_correct += (net(X).argmax(1) == y).sum().item()
        test_acc = test_correct / len(test_iter.dataset)

        # 记录结果
        train_loss_list.append(train_loss)
        train_acc_list.append(train_acc)
        test_acc_list.append(test_acc)

        print(f'Epoch {epoch + 1}: '
              f'Train Loss: {train_loss:.4f}, Acc: {train_acc:.3f} | '
              f'Test Acc: {test_acc:.3f}')

    # 简洁可视化（如您要求的格式）
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(train_loss_list, label='Train Loss')
    plt.xlabel('Epoch')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(train_acc_list, label='Train Acc')
    plt.plot(test_acc_list, label='Test Acc')
    plt.xlabel('Epoch')
    plt.legend()

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    torch.manual_seed(42)

    # 参数设置
    batch_size = 256
    num_epochs = 10
    lr = 0.1
    weight_decay = 1e-4  # L2正则化系数
    dropout_rate = 0.5  # Dropout率

    # 加载数据
    train_iter, test_iter = load_data_mnist(batch_size)

    # 初始化模型（带Dropout正则化）
    model = RegularizedNN(dropout_rate=dropout_rate)

    # 训练模型（带L2正则化）
    train_model(model, train_iter, test_iter,
                num_epochs=num_epochs,
                lr=lr,
                weight_decay=weight_decay)